﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Epsilon.Expr;

namespace Epsilon
{
    public static class SecUtil
    {
        private static int sec = 0;

        private static void Error(string e)
        {
            throw new SecLabelException(e);
        }

        public static SecVar NewVar()
        {
            sec++;
            return new SecVar("_S" + sec.ToString());
        }

        public static SecFun AlphaConvert(SecFun f, Dictionary<SecVar, SecVar> sub)
        {
            List<SecRLabel> args = new List<SecRLabel>();
            for (int i = 0; i < f.Args.Count; i++)
                args.Add(Alpha(sub, f.Args[i]));

            Dictionary<SecVar, SecBase> newSub = new Dictionary<SecVar, SecBase>();
            foreach (SecVar v in sub.Keys)
            {
                List<SecVar> s = new List<SecVar>();
                s.Add(sub[v]);
                newSub.Add(v, new SecBase(s, null));
            }

            SecLabel ret = Sub(newSub, f.Ret);
            return new SecFun(args, ret);
        }

        #region Meet

        public static SecLabel Meet(List<SecLabel> labels)
        {
            if (labels.Count == 0)
                return new SecBase(null, null);

            SecLabel p = labels[0];
            for (int i = 1; i < labels.Count; i++)
                p = Meet(p, labels[i]);

            return p;
        }

        public static SecLabel Meet(SecLabel left, SecLabel right)
        {
            if ((left is SecBase) && ((SecBase)left).IsTop)
                return right;
            else if ((right is SecBase) && ((SecBase)right).IsTop)
                return left;

            if ((left is SecBase) && (right is SecBase))
                return Meet((SecBase)left, (SecBase)right);
            else if ((left is SecTuple) && (right is SecTuple))
                return Meet((SecTuple)left, (SecTuple)right);
            else if ((left is SecFun) && (right is SecFun))
                return Meet((SecFun)left, (SecFun)right);
            Error("Unable to compute the meet of " + left.ToString() +
                " and " + right.ToString());
            return null;
        }

        public static SecBase Meet(SecBase left, SecBase right)
        {
            if (left.IsTop)
                return right;
            else if (right.IsTop)
                return left;

            List<SecVar> secvs = null;
            List<SecPolicy> policies = null;

            if (left.Secvs != null)
            {
                secvs = new List<SecVar>();
                foreach (SecVar v in left.Secvs)
                    secvs.Add(v);
            }
            if (left.Policies != null)
            {
                policies = new List<SecPolicy>();
                foreach (SecPolicy p in left.Policies)
                    policies.Add(p);
            }

            if (right.Secvs != null)
            {
                if (secvs == null)
                    secvs = right.Secvs;
                else
                    secvs = new List<SecVar>(secvs.Concat(right.Secvs).Distinct());
            }
            if (right.Policies != null)
            {
                if (policies == null)
                    policies = right.Policies;
                else
                    policies = MeetPolicies(left.Policies, right.Policies);
            }

            return new SecBase(secvs, policies);
        }

        public static SecTuple Meet(SecTuple left, SecTuple right)
        {
            if (left.Labels.Count != right.Labels.Count)
                Error("Unable to meet tuples of different lengths: "
                    + left.ToString() + " and " + right.ToString());

            List<SecLabel> labels = new List<SecLabel>();
            for (int i = 0; i < left.Labels.Count; i++)
                labels.Add(Meet(left.Labels[i], right.Labels[i]));
            return new SecTuple(labels);
        }

        public static SecFun Meet(SecFun left, SecFun right)
        {
            // check for alpha conversion
            SecRTuple largs = new SecRTuple(left.Args);
            SecRTuple rargs = new SecRTuple(right.Args);

            List<SecRLabel> args = left.Args;
            SecLabel lBody = left.Ret;
            SecLabel rBody = right.Ret;

            if (largs != rargs)
            {
                args = new List<SecRLabel>();
                foreach (SecRLabel a in left.Args)
                    args.Add(CreateNewAlpha(a));
                // Create a sec tuple for ease of calculating the substitution
                SecRTuple lsrc = new SecRTuple(left.Args);
                SecRTuple rsrc = new SecRTuple(right.Args);
                SecRTuple dest = new SecRTuple(args);
                Dictionary<SecVar, SecVar> lsub = Equate(lsrc, dest);
                Dictionary<SecVar, SecVar> rsub = Equate(rsrc, dest);
                lBody = AlphaConvert(left, lsub).Ret;
                rBody = AlphaConvert(right, rsub).Ret;
            }

            SecLabel ret = Meet(lBody, rBody);
            return new SecFun(args, ret);
        }

        #endregion

        #region Vars 

        public static List<SecVar> Vars(SecLabel l)
        {
            if (l is SecBase)
                return Vars((SecBase)l);
            else if (l is SecTuple)
                return Vars((SecTuple)l);
            else if (l is SecFun)
                return Vars((SecFun)l);
            else
                Error("Invalid security label");
            return null;
        }

        public static List<SecVar> Vars(SecBase b)
        {
            return b.Secvs;
        }

        public static List<SecVar> Vars(SecTuple t)
        {
            List<SecVar> v = new List<SecVar>();
            for (int i = 0; i < t.Labels.Count; i++)
                v.AddRange(Vars(t.Labels[i]));
            return new List<SecVar>(v.Distinct());
        }

        public static List<SecVar> Vars(SecFun f)
        {
            List<SecVar> v = new List<SecVar>();
            for (int i = 0; i < f.Args.Count; i++)
                v.AddRange(Vars(f.Args[i]));
            v.AddRange(Vars(f.Ret));
            return new List<SecVar>(v.Distinct());
        }

        public static List<SecVar> Vars(SecRLabel l)
        {
            if (l is SecVar)
                return Vars((SecVar)l);
            else if (l is SecRTuple)
                return Vars((SecRTuple)l);
            else if (l is SecRFun)
                return Vars((SecRFun)l);
            else
                Error("Invalid restricted label " + l.ToString());
            return null;
        }

        public static List<SecVar> Vars(SecVar v)
        {
            List<SecVar> l = new List<SecVar>();
            l.Add(v);
            return l;
        }

        public static List<SecVar> Vars(SecRTuple t)
        {
            List<SecVar> v = new List<SecVar>();
            for (int i = 0; i < t.Labels.Count; i++)
                v.AddRange(Vars(t.Labels[i]));
            return new List<SecVar>(v.Distinct());
        }

        public static List<SecVar> Vars(SecRFun f)
        {
            List<SecVar> v = new List<SecVar>();
            for (int i = 0; i < f.Args.Count; i++)
                v.AddRange(Vars(f.Args[i]));
            v.AddRange(Vars(f.Ret));
            return new List<SecVar>(v.Distinct());
        }

        #endregion

        #region CreateNewAlpha

        public static SecRLabel CreateNewAlpha(SecRLabel l)
        {
            if (l is SecVar)
                return CreateNewAlpha((SecVar)l);
            else if (l is SecRTuple)
                return CreateNewAlpha((SecRTuple)l);
            else if (l is SecRFun)
                return CreateNewAlpha((SecRFun)l);
            else
                Error("Invalid restricted label " + l.ToString());
            return null;
        }

        public static SecRLabel CreateNewAlpha(SecVar v)
        {
            return NewVar();
        }

        public static SecRLabel CreateNewAlpha(SecRTuple t)
        {
            List<SecRLabel> labels = new List<SecRLabel>();
            for (int i = 0; i < t.Labels.Count; i++)
                labels.Add(CreateNewAlpha(t.Labels[i]));
            return new SecRTuple(labels);
        }

        public static SecRLabel CreateNewAlpha(SecRFun f)
        {
            List<SecRLabel> args = new List<SecRLabel>();
            for (int i = 0; i < f.Args.Count; i++)
                args.Add(CreateNewAlpha(f.Args[i]));
            return new SecRFun(args, CreateNewAlpha(f.Ret));
        }

        #endregion

        #region Equate

        public static Dictionary<SecVar, SecVar> Equate(SecRLabel src, SecRLabel dest)
        {
            if ((src is SecVar) && (dest is SecVar))
                return Equate((SecVar)src, (SecVar)dest);
            else if ((src is SecRTuple) && (dest is SecRTuple))
                return Equate((SecRTuple)src, (SecRTuple)dest);
            else if ((src is SecRFun) && (dest is SecRFun))
                return Equate((SecRFun)src, (SecRFun)dest);
            Error("Unable to equate " + src.ToString() +
                " and " + dest.ToString());
            return null;
        }

        public static Dictionary<SecVar, SecVar> Equate(SecVar src, SecVar dest)
        {
            Dictionary<SecVar, SecVar> sub = new Dictionary<SecVar, SecVar>();
            sub.Add(src, dest);

            return sub;
        }

        public static Dictionary<SecVar, SecVar> Equate(SecRTuple src, SecRTuple dest)
        {
            if (src.Labels.Count != dest.Labels.Count)
                Error("Unable to equate " + src.ToString() + " and " + dest.ToString());

            Dictionary<SecVar, SecVar> sub = new Dictionary<SecVar, SecVar>();

            for (int i = 0; i < src.Labels.Count; i++)
            {
                Dictionary<SecVar, SecVar> s = Equate(src.Labels[i], dest.Labels[i]);
                foreach (SecVar v in s.Keys)
                {
                    if (sub.ContainsKey(v))
                        Error("Duplicate variable in " + src.ToString());
                    sub.Add(v, s[v]);
                }
            }

            return sub;
        }

        public static Dictionary<SecVar, SecVar> Equate(SecRFun src, SecRFun dest)
        {
            if (src.Args.Count != dest.Args.Count)
                Error("Unable to equate " + src.ToString() + " and " + dest.ToString());

            Dictionary<SecVar, SecVar> sub = new Dictionary<SecVar, SecVar>();

            for (int i = 0; i < src.Args.Count; i++)
            {
                Dictionary<SecVar, SecVar> s = Equate(src.Args[i], dest.Args[i]);
                foreach (SecVar v in s.Keys)
                {
                    if (sub.ContainsKey(v))
                        Error("Duplicate variable in " + src.ToString());
                    sub.Add(v, s[v]);
                }
            }

            Dictionary<SecVar, SecVar> ss = Equate(src.Ret, dest.Ret);
            foreach (SecVar v in ss.Keys)
            {
                if (sub.ContainsKey(v))
                    Error("Duplicate variable in " + src.ToString());
                sub.Add(v, ss[v]);
            }

            return sub;
        }

        #endregion

        #region Sub

        public static SecLabel Sub(Dictionary<SecVar, SecBase> sub, SecLabel label)
        {
            if (label is SecBase)
                return Sub(sub, (SecBase)label);
            else if (label is SecTuple)
                return Sub(sub, (SecTuple)label);
            else if (label is SecFun)
                return Sub(sub, (SecFun)label);
            else
                Error("Invalid security label " + label.ToString());
            return null;
        }

        public static SecBase Sub(Dictionary<SecVar, SecBase> sub, SecBase label)
        {
            if (label.Secvs == null || label.Secvs.Count == 0)
                return label;

            List<List<SecPolicy>> policies = new List<List<SecPolicy>>();
            List<List<SecVar>> vars = new List<List<SecVar>>();

            policies.Add(label.Policies);
            foreach (SecVar v in label.Secvs)
            {
                if (sub.ContainsKey(v))
                {
                    policies.Add(sub[v].Policies);
                    vars.Add(sub[v].Secvs);
                }
                else
                {
                    vars.Add(new List<SecVar>() { v });
                }
            }

            policies.RemoveAll(item => item == null);
            vars.RemoveAll(item => item == null);

            List<SecPolicy> newPolicies = null;
            IEnumerable<SecVar> newVars = null;
            if (policies.Count > 0)
            {
                newPolicies = policies[0];
                for (int i = 1; i < policies.Count; i++)
                    newPolicies = MeetPolicies(newPolicies, policies[i]);
            }
            if (vars.Count > 0)
            {
                newVars = vars[0];
                for (int i = 1; i < vars.Count; i++)
                    newVars = newVars.Union(vars[i]);
            }

            List<SecVar> nv = (newVars == null) ? null : new List<SecVar>(newVars);
            return new SecBase(nv, newPolicies);
        }

        public static SecTuple Sub(Dictionary<SecVar, SecBase> sub, SecTuple label)
        {
            List<SecLabel> labels = new List<SecLabel>();
            for (int i = 0; i < label.Labels.Count; i++)
                labels.Add(Sub(sub, label.Labels[i]));
            return new SecTuple(labels);
        }

        public static SecFun Sub(Dictionary<SecVar, SecBase> sub, SecFun label)
        {
            bool alphaConvert = false;

            List<SecVar> subVars = new List<SecVar>();
            foreach (SecVar v in sub.Keys)
            {
                subVars.Add(v);
                subVars.AddRange(sub[v].Secvs);
            }

            foreach (SecRLabel a in label.Args)
            {
                if (Vars(a).Intersect(subVars).Count() > 0)
                    alphaConvert = true;
            }

            SecFun res = label;

            if (alphaConvert)
            {
                List<SecRLabel> newArgs = new List<SecRLabel>();
                foreach (SecRLabel a in label.Args)
                    newArgs.Add(CreateNewAlpha(a));
                // Create a Sec tuple for ease of calculating the substitution
                SecRTuple src = new SecRTuple(label.Args);
                SecRTuple dest = new SecRTuple(newArgs);
                Dictionary<SecVar, SecVar> asub = Equate(src, dest);
                res = AlphaConvert(label, asub);
            }

            // Now return the final function
            List<SecRLabel> resArgs = res.Args;
            SecLabel resRet = Sub(sub, res.Ret);
            return new SecFun(resArgs, resRet);
        }

        #endregion

        #region Alpha

        public static SecRLabel Alpha(Dictionary<SecVar, SecVar> sub, SecRLabel label)
        {
            if (label is SecVar)
                return Alpha(sub, (SecVar)label);
            else if (label is SecRTuple)
                return Alpha(sub, (SecRTuple)label);
            else if (label is SecRFun)
                return Alpha(sub, (SecRFun)label);
            else
                Error("Invalid restricted label " + label.ToString());
            return null;
        }

        public static SecRLabel Alpha(Dictionary<SecVar, SecVar> sub, SecVar label)
        {
            if (sub.ContainsKey(label))
                return sub[label];
            else
                return label;
        }

        public static SecRLabel Alpha(Dictionary<SecVar, SecVar> sub, SecRTuple label)
        {
            List<SecRLabel> labels = new List<SecRLabel>();
            for (int i = 0; i < label.Labels.Count; i++)
                labels.Add(Alpha(sub, label.Labels[i]));
            return new SecRTuple(labels);
        }

        public static SecRLabel Alpha(Dictionary<SecVar, SecVar> sub, SecRFun label)
        {
            List<SecRLabel> args = new List<SecRLabel>();
            for (int i = 0; i < label.Args.Count; i++)
                args.Add(Alpha(sub, label.Args[i]));
            return new SecRFun(args, Alpha(sub, label.Ret));
        }

        #endregion

        #region MeetPolicies

        public static List<SecPolicy> MeetPolicies(List<SecPolicy> left, List<SecPolicy> right)
        {
            List<SecPolicy> policies = new List<SecPolicy>();

            // first add all policies with an owner not present in the other list
            foreach (SecPolicy sp in left)
            {
                bool found = false;
                foreach (SecPolicy ssp in right)
                {
                    if (EqualOwners(sp, ssp))
                    {
                        found = true;
                        break;
                    }
                }
                if (!found)
                    policies.Add(sp);
            }

            foreach (SecPolicy sp in right)
            {
                bool found = false;
                SecPolicy nsp = null;
                foreach (SecPolicy ssp in left)
                {
                    if (EqualOwners(sp, ssp))
                    {
                        found = true;
                        nsp = ssp;
                        break;
                    }
                }
                if (!found)
                    policies.Add(sp);
                // otherwise calculate the meet and add that
                else
                    policies.Add(MeetPol(sp, nsp));
            }

            return policies;
        }

        public static bool EqualOwners(SecPolicy l, SecPolicy r)
        {
            return EqualSecPids(l.Owner, r.Owner);
        }

        public static bool EqualSecPids(ISecPid left, ISecPid right)
        {
            if ((left is PidVar) && (right is PidVar))
                return ((PidVar)left).PidV == ((PidVar)right).PidV;
            else if ((left is PidExpr) && (right is PidExpr))
                return ((PidExpr)left).Pid == ((PidExpr)right).Pid;

            return false;
        }

        public static SecPolicy MeetPol(SecPolicy l, SecPolicy r)
        {
            if (!EqualOwners(l, r))
                Error("Unable to meet two policies with different owners");

            List<ISecPid> secpids = new List<ISecPid>();

            foreach (ISecPid p in l.Readers)
            {
                bool found = false;
                foreach (ISecPid pp in r.Readers)
                {
                    if (EqualSecPids(p, pp))
                    {
                        found = true;
                        break;
                    }
                }
                if (found)
                    secpids.Add(p);
            }

            return new SecPolicy(l.Owner, secpids);
        }

        #endregion

        #region Taint

        public static SecLabel Taint(SecLabel src, SecLabel taint)
        {
            if (src is SecBase)
                return Taint((SecBase)src, taint);
            else if (src is SecTuple)
                return Taint((SecTuple)src, taint);
            else if (src is SecFun)
                return Taint((SecFun)src, taint);
            else
                Error("Invalid security label " + src.ToString());
            return null;
        }

        public static SecBase Taint(SecBase src, SecLabel taint)
        {
            if (taint is SecBase)
                return Taint((SecBase)src, (SecBase)taint);
            else if (taint is SecTuple)
                return Taint((SecBase)src, (SecTuple)taint);
            else if (taint is SecFun)
                return Taint((SecBase)src, (SecFun)taint);
            else
                Error("Invalid security label " + taint.ToString());
            return null;
        }

        public static SecBase Taint(SecBase src, SecBase taint)
        {
            return Meet(src, taint);
        }

        public static SecBase Taint(SecBase src, SecTuple taint)
        {
            if (taint.Labels.Count == 0)
                Error("Unable to taint " + src.ToString() + " with empty tuple");

            SecBase ret = Taint(src, taint.Labels[0]);
            for (int i = 1; i < taint.Labels.Count; i++)
                ret = Meet(ret, Taint(src, taint.Labels[i]));
            return ret;
        }

        public static SecBase Taint(SecBase src, SecFun taint)
        {
            return Taint(src, taint.Ret);
        }

        public static SecTuple Taint(SecTuple src, SecLabel taint)
        {
            List<SecLabel> labels = new List<SecLabel>();
            for (int i = 0; i < src.Labels.Count; i++)
                labels.Add(Taint(src.Labels[i], taint));
            return new SecTuple(labels);
        }

        public static SecFun Taint(SecFun src, SecLabel taint)
        {
            return new SecFun(src.Args, Taint(src.Ret, taint));
        }

        #endregion

        #region CreateFromType

        public static SecRLabel CreateFromType(Type t)
        {
            if (t is SimpleType)
                return CreateFromType((SimpleType)t);
            else if (t is AtomType)
                return CreateFromType((AtomType)t);
            else if (t is ListType)
                return CreateFromType((ListType)t);
            else if (t is TupleType)
                return CreateFromType((TupleType)t);
            else if (t is FunType)
                return CreateFromType((FunType)t);
            else
                Error("Invalid type");
            return null;
        }

        public static SecVar CreateFromType(SimpleType t)
        {
            return NewVar();
        }

        public static SecVar CreateFromType(AtomType t)
        {
            return NewVar();
        }

        public static SecRLabel CreateFromType(ListType t)
        {
            return CreateFromType(t.ElementType);
        }

        public static SecRTuple CreateFromType(TupleType t)
        {
            List<SecRLabel> labels = new List<SecRLabel>();
            for (int i = 0; i < t.Types.Count; i++)
                labels.Add(CreateFromType(t.Types[i]));
            return new SecRTuple(labels);
        }

        public static SecRFun CreateFromType(FunType t)
        {
            List<SecRLabel> args = new List<SecRLabel>();
            for (int i = 0; i < t.Args.Count; i++)
                args.Add(CreateFromType(t.Args[i]));
            return new SecRFun(args, CreateFromType(t.Ret));
        }

        #endregion

        #region FromRestricted

        public static SecLabel FromRestricted(SecRLabel src)
        {
            if (src is SecVar)
                return FromRestricted((SecVar)src);
            else if (src is SecRTuple)
                return FromRestricted((SecRTuple)src);
            else if (src is SecRFun)
                return FromRestricted((SecRFun)src);
            else
                Error("Invalid restricted security label");
            return null;
        }

        public static SecLabel FromRestricted(SecVar src)
        {
            return new SecBase(new List<SecVar>() { src }, null);
        }

        public static SecLabel FromRestricted(SecRTuple src)
        {
            List<SecLabel> labels = new List<SecLabel>();
            for (int i = 0; i < src.Labels.Count; i++)
                labels.Add(FromRestricted(src.Labels[i]));
            return new SecTuple(labels);
        }

        public static SecLabel FromRestricted(SecRFun src)
        {
            return new SecFun(src.Args, FromRestricted(src.Ret));
        }

        #endregion

        #region Unify

        public static Dictionary<SecVar, SecBase> Unify(SecRLabel src, SecLabel dest)
        {
            if (src is SecVar)
                if (dest is SecVar)
                    return Unify((SecVar)src, (SecVar)dest);
                else if (dest is SecBase)
                    return Unify((SecVar)src, (SecBase)dest);
                else
                    Error("Unable to unify " + src.ToString() + " and " + dest.ToString());
            else if ((src is SecRTuple) && (dest is SecTuple))
                return Unify((SecRTuple)src, (SecTuple)dest);
            else if ((src is SecRFun) && (dest is SecFun))
                return Unify((SecRFun)src, (SecFun)dest);
            else
                Error("Unable to unify " + src.ToString() + " and " + dest.ToString());
            return null;
        }

        public static Dictionary<SecVar, SecBase> Unify(SecVar src, SecVar dest)
        {
            Dictionary<SecVar, SecBase> dict = new Dictionary<SecVar, SecBase>();
            dict.Add(src, new SecBase(new List<SecVar>() { dest }, null));
            return dict;
        }

        public static Dictionary<SecVar, SecBase> Unify(SecVar src, SecBase dest)
        {
            Dictionary<SecVar, SecBase> dict = new Dictionary<SecVar, SecBase>();
            dict.Add(src, dest);
            return dict;
        }

        public static Dictionary<SecVar, SecBase> Unify(SecRTuple src, SecTuple dest)
        {
            if (src.Labels.Count != dest.Labels.Count)
                Error("Unable to unify " + src.ToString() + " and " + dest.ToString());

            Dictionary<SecVar, SecBase> dict = new Dictionary<SecVar, SecBase>();

            for (int i = 0; i < src.Labels.Count; i++)
            {
                Dictionary<SecVar, SecBase> sub = Unify(src.Labels[i], dest.Labels[i]);
                foreach (SecVar v in sub.Keys)
                {
                    if (dict.ContainsKey(v))
                        Error("Duplicate variable in " + src.ToString());
                    dict.Add(v, sub[v]);
                }
            }

            return dict;
        }

        public static Dictionary<SecVar, SecBase> Unify(SecRFun src, SecFun dest)
        {
            if (src.Args.Count != dest.Args.Count)
                Error("Unable to unify " + src.ToString() + " and " + dest.ToString());

            Dictionary<SecVar, SecBase> dict = new Dictionary<SecVar, SecBase>();

            for (int i = 0; i < src.Args.Count; i++)
            {
                Dictionary<SecVar, SecBase> sub = Unify(src.Args[i], dest.Args[i]);
                foreach (SecVar v in sub.Keys)
                {
                    if (dict.ContainsKey(v))
                        Error("Duplicate variable in " + src.ToString());
                    dict.Add(v, sub[v]);
                }
            }

            Dictionary<SecVar, SecBase> ssub = Unify(src.Ret, dest.Ret);
            foreach (SecVar v in ssub.Keys)
            {
                if (dict.ContainsKey(v))
                    Error("Duplicate variable in " + src.ToString());
                dict.Add(v, ssub[v]);
            }

            return dict;
        }

        #endregion
    }
}
